from langchain.agents import Tool, AgentExecutor
from langchain.agents.agent_types import AgentType
from langchain_LCEL import get_LECL_chain
from langchain_LCEL import get_llm
from langchain_core.prompts import PromptTemplate,FewShotPromptTemplate,SystemMessagePromptTemplate
from langgraph.prebuilt import create_react_agent
from langchain.schema.runnable import RunnableMap,RunnableParallel,RunnableBranch,RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain.schema import HumanMessage
import json

registry={}
# 装饰器的原理，其实就是 “函数接收函数，并返回函数” —— 本质是 高阶函数 和 闭包 的组合应用。
"""
装饰器本质上是一个“函数”，它接收一个函数作为参数，加工后再返回一个新函数（可以是原函数，也可以是包裹的新函数），用于增强函数的功能。
"""
def register(name):
    def decorator(func):
        registry[name] = func
        return func
    return decorator


def paser_history_str(input_dict):
    # raw_history structured_summary
    prompt_input= {}
    try:
        raw_history = input_dict["raw_history"]
        structured_summary =  input_dict["structured_summary"]
        prompt_input = {
        "raw_history": raw_history,
        "structured_summary": structured_summary
        }
    except Exception as e:
        print(e)
    # 构造符合 FewShotPromptTemplate 的输入格式
    if prompt_input: 
        return prompt_input
    else:
        None
        
def paser_parse_str(input_dict):
    prompt_input = None
    # raw_history structured_summary
    try:
        symptom = input_dict["symptom"]
        structured_summary = input_dict['structured_summary']
        prompt_input = {
        "symptom": symptom,
        "structured_summary": structured_summary
        }
    
    except Exception as e:
        print(e)
    if prompt_input: 
        return prompt_input
    else:
        None

def paser_diagnosis_str(input_str):
    if isinstance(input_str,dict):
        return input_str
    # raw_history structured_summary
    parsed = json.loads(input_str)
    symptom = parsed.get("symptom", "无症状")
    structured_summary = parsed.get("structured_summary", "无总结病史")
    diagnosis_chat_history = parsed.get("diagnosis_chat_history", "无问答历史")
    # 构造符合 FewShotPromptTemplate 的输入格式
    prompt_input = {
        "symptom": symptom,
        "structured_summary": structured_summary,
        "diagnosis_chat_history":diagnosis_chat_history
    }
    return prompt_input

def paser_recipe_str(input_str):
    # raw_history structured_summary
    parsed = json.loads(input_str)
    symptom = parsed.get("symptom", "无症状")
    structured_summary = parsed.get("structured_summary", "无总结病史")
    diagnosis_chat_history = parsed.get("diagnosis_chat_history", "无问答历史")
    # 构造符合 FewShotPromptTemplate 的输入格式
    prompt_input = {
        "symptom": symptom,
        "structured_summary": structured_summary,
        "diagnosis_chat_history":diagnosis_chat_history
    }
    return prompt_input

def paser_recheck_str(input_str):
    # raw_history structured_summary
    parsed = json.loads(input_str)
    recipe = parsed.get("recipe", "无药方")
    structured_summary = parsed.get("structured_summary", "无总结病史")
    diagnosis_chat_history = parsed.get("diagnosis_chat_history", "无问答历史")
    # 构造符合 FewShotPromptTemplate 的输入格式
    prompt_input = {
        "recipe": recipe,
        "structured_summary": structured_summary,
        "diagnosis_chat_history":diagnosis_chat_history
    }
    return prompt_input


@register('history_agent')
def build_history_agent():
   # 单个示例格式
    example_prompt = PromptTemplate(
    input_variables=["raw_history", "structured_summary"],
    template="""
    【原始就诊记录】
    {raw_history}

    【已有总结就诊记录】
    {structured_summary}
    """.strip()
    )

    examples = [
        {
            "raw_history": """2023年5月10日
    患者主诉：反复咳嗽两周，痰白，清晨明显。
    医生：有发热吗？有没有鼻塞流涕？
    患者：没有发热，有点鼻塞。
    医生：有没有既往过敏史？
    患者：有慢性鼻炎。
    医生判断：考虑风寒袭肺，伴有痰湿。
    建议：开二陈汤加苏叶，避免受凉。

    2023年6月5日
    患者：最近咳嗽又加重，晚上也咳，影响睡眠。
    医生：是否仍为白痰？有没有咽喉不适？
    患者：痰粘稠，有点黄，喉咙干痒。
    医生判断：风寒转化为痰热犯肺。
    建议：改用清金化痰汤，嘱咐忌食辛辣。""",
            "structured_summary": """
    【2023年5月10日】
    - 主诉：反复咳嗽两周，清晨痰白
    - 病史要点：慢性鼻炎史；无发热，有鼻塞
    - 判断：风寒袭肺，痰湿内生
    - 处理建议：二陈汤加苏叶，嘱避免受凉

    【2023年6月5日】
    - 主诉：咳嗽加重，夜间咳嗽，痰粘稠微黄
    - 病史要点：咽干痒，睡眠受影响
    - 判断：风寒化热，痰热犯肺
    - 处理建议：清金化痰汤，忌食辛辣
    """
        },
        {
            "raw_history": "无就诊记录",
            "structured_summary": "无就诊记录"
        }
    ]

    prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
        prefix="""
    你是中医问诊记录分析专家。请按如下规则处理任务：

    1. 如果已提供总结（structured_summary），请比对原始就诊记录（raw_history），检查是否有遗漏或错误信息，并在“补充信息”中补全。
    2. 如果未提供 structured_summary，请你从 raw_history 中提取结构化的病史总结。
    3. 如果 raw_history 是“无就诊记录”，直接输出“无就诊信息”。

    输出格式如下（根据情况选择）：
    - 若需补全，请列出“补充信息”
    - 若是新总结，请输出结构化总结
    - 若无记录，请输出：无就诊信息
    """.strip(),
        suffix="""
    【原始就诊记录】
    {raw_history}

    【已有总结就诊记录】
    {structured_summary}

    【输出】
    """.strip(),
        input_variables=["raw_history", "structured_summary"]
    )
    # RunnableLambda(lambda x,{})
    output_parser = StrOutputParser()
    llm = get_llm()
    history_lecl_chain = RunnableLambda(paser_history_str)|prompt|llm|output_parser
    # history_lecl_chain.invoke("")
    
    # 包装为 Tool
    # my_tool = Tool(
    #     name="build_history_tool",
    #     func=history_lecl_chain.invoke,  
    #     description=(
    #         """这个工具可以将病历和已经有的病历总结进行对比，返回再次总结后的病史。
    #             工具要求的输入格式为："raw_history":"病人问诊记录","structured_summary":"历史总结就诊记录"
    #         """
    #     )
    # )
    # """
    # create_react_agent() 用的是 ReAct agent，它默认使用 messages、intermediate_steps、input 等格式来运行。
    # 你不能直接把一个专门用于任务型 Prompt（如你现在用于“提取病史”的 FewShotPromptTemplate）当作 ReAct Agent 的整体 Prompt。
    # """
    
    
    # agent = create_react_agent(
    #     model=llm,
    #     tools=[my_tool],
    #     prompt="""你会得到一段病人的历史问诊记录和已经有的病史总结文本，请你调用已有的工具对病人病史进行总结。
    #     如果需要调用工具，请按照工具预期的输入格式将文本转换成工具预期输入格式。""" # 可选，自定义 ReAct Prompt
    # )
    return history_lecl_chain

@register('parse_agent')
def build_parse_agent():
    prompt = PromptTemplate(
        input_variables=["symptom", "structured_summary"],
        template="""
        请根据病人总结病史和病人当前症状，返回更加专业的症状描述。请在你的回答之前添加，患者主诉。
        
        病人病史：
        {structured_summary}
        
        病人症状：
        {symptom}
        
        你的症状描述：
        患者主诉：
    """.strip()
    )
    
    
    # system_message_prompt = SystemMessagePromptTemplate.from_template(prompt)
    output_parser = StrOutputParser()
    llm = get_llm()
    
    parse_lecl_chain = RunnableLambda(paser_parse_str)|prompt|llm|output_parser
    # res = parse_lecl_chain.invoke({"symptom":"有点不舒服","structured_summary":"无总结病史"})
    # print(res)
    # 包装为 Tool
    
    # my_tool = Tool(
    #     name="build_parse_tool",
    #     func=parse_lecl_chain.invoke,  
    #     description=
    #     """
    #     这个工具可以将病人症状和已经有的病史总结进行结合，返回更加专业的患者病情描述。
    #     工具要求的输入格式为："symptom":"病人症状","structured_summary":"病人病史"
    #     """
    # )
    # # res = my_tool.invoke('{"symptom":"有点不舒服","structured_summary":"无总结病史"}')
    # # print(res)
    # agent = create_react_agent(
    #     model=llm,
    #     tools=[my_tool],
    #     prompt="""你会得到一段病人的病史总结和病人自述症状，请你调用已有的工具进行处理，并返回工具的结果。如果需要调用工具，请按照工具预期的输入格式将文本转换成工具预期输入格式。
    #     如果患者不是在叙述病情，请提醒他叙述病情，语气需要柔和。""" # 可选，自定义 ReAct Prompt
    # )
    
    # return agent
    # parse_lecl_chain = RunnableLambda(paser_parse_str)|prompt|llm|output_parser
    return  parse_lecl_chain
    
@register('diagnosis_agent')
def build_diagnosis_agent():
    example_prompt = PromptTemplate(
        input_variables=["structured_summary", "symptom", "diagnosis_chat_history"],
        template="""
    【中医逐步问诊流程】
    根据当前信息，请按步骤执行问诊（每次只提1个最关键问题）：

    ▼▼▼ 当前已知信息 ▼▼▼
    [既往病史]
    {structured_summary}

    [主诉症状]
    {symptom}

    [历史对话]
    {diagnosis_chat_history}

    ▲▲▲ 请执行下一步 ▲▲▲
    请严格按以下逻辑执行：
    1. 如果对话中未出现"过敏史"关键词，必须首先询问：
    → "请问您是否有药物、食物或中药材过敏史？"
    
    2. 如果已确认过敏史但症状信息不足，根据主诉选择最关键的1个问题：
    - 失眠相关：询问作息时间/情绪状态
    - 咳嗽相关：询问痰液性质/持续时间
    - 疼痛相关：询问疼痛性质/诱发因素
    （其他症状类推）

    3. 如果已获取完整信息，则输出：
    【中医辨证】
    [病机分析]
    
    【治疗方案】
    [具体建议]
    """.strip()
    )

    examples = [
        {
            "structured_summary": "高血压病史3年",
            "symptom": "失眠多梦伴口苦1个月",
            "diagnosis_chat_history": "",
            "output": "请问您是否有药物、食物或中药材过敏史？"
        },
        {
            "structured_summary": "无",
            "symptom": "反复头痛3个月",
            "diagnosis_chat_history": "医生：是否有过敏史？\n患者：对青霉素过敏",
            "output": "头痛的具体部位在哪里？（如太阳穴/头顶/后脑）"
        },
        {
            "structured_summary": "吸烟史10年",
            "symptom": "晨起咳嗽痰多2周",
            "diagnosis_chat_history": "医生：是否有过敏史？\n患者：无\n医生：痰液是什么颜色？\n患者：黄色粘稠",
            "output": """
    【中医辨证】
    证属痰热壅肺。烟毒灼津成痰，痰热互结，肺失宣降。

    【治疗方案】
    • 方药：清金化痰汤加减（黄芩10g、瓜蒌15g、浙贝母10g）
    • 针灸：肺俞、尺泽、丰隆
    • 饮食：忌辛辣油腻，推荐雪梨百合羹
    """
        }
    ]

    prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix="""作为中医专家，您需要：
1. 严格遵循"过敏史→核心症状→伴随症状"的问诊顺序
2. 每次只提1个最关键问题
3. 确保问题之间有逻辑递进关系""",
    suffix="""
▼▼▼ 当前就诊信息 ▼▼▼
[主诉] {symptom}
[病史] {structured_summary}
[对话记录] {diagnosis_chat_history}

请输出您的下一步问诊或诊断：""",
    input_variables=["symptom", "structured_summary", "diagnosis_chat_history"]
)
    
    output_parser = StrOutputParser()
    llm = get_llm()
    diagnosis_lecl_chain = RunnableLambda(paser_diagnosis_str)|prompt|llm|output_parser
    # 包装为 Tool
    my_tool = Tool(
        name="build_diagnosis_tool",
        func=diagnosis_lecl_chain.invoke,  
        description="该工具进行病情诊断"
    )
    agent = create_react_agent(
        model=llm,
        tools=[my_tool],
        prompt="你需要根据输入的文本，判断是否需要工具进行操作。工具预期的输入格式为'structured_summary':'...'和'symptom':'...'以及'diagnosis_chat_history':'...'的字典。如果患者不是在叙述病情，请提醒他叙述病情，语气需要柔和。判断工具的返回结果是否正确，如果正确即刻返回结果。" # 可选，自定义 ReAct Prompt
    )
    return diagnosis_lecl_chain
    # return agent

@register('recipe_agent')
def build_recipe_agent():
    recipe_lecl_chain = get_LECL_chain()
    refine_tool = Tool(
        name=" build_recipe_tool",
        func=recipe_lecl_chain.invoke,
        description="使用该链给出药方"
    )
    llm = get_llm()
    agent = create_react_agent(
        tools=[refine_tool],  # 可支持多个 tool
        model=llm,
        prompt="你需要根据输入的文本，判断是否需要工具进行操作。工具预期的输入格式为'diagnosis_chat_history':'...'的字典。判断工具的返回结果是否正确，如果正确即刻返回结果。"
    )
    return recipe_lecl_chain
    # return agent

@register('recheck_agent')   
def build_recheck_agent():
    prompt = PromptTemplate.from_template(
        """
        你现在是一名专业的中医药医生，你需要将给出的药方与病人病史相结合之后的诊断历史对话，判断这个药方是否合适，并给出你的理由。如果没有病人病史则结合诊断判断是否合适。如果这个药方是合适的，在答案处返回，完全合适。
        病人病史：
        {structured_summary}
        诊断对话：
        {diagnosis_chat_history}
        药方：
        {recipe}
        你的回答请按照以下格式：
        答案：你的回答
        理由：你的理由
        # """
        # input_variables=["recipe", "structured_summary", "diagnosis_chat_history"]
    )
    output_parser = StrOutputParser()
    llm = get_llm()
    # recheck_lecl_chain = RunnableLambda(paser_recheck_str)|prompt|llm|output_parser
    recheck_lecl_chain = prompt|llm|output_parser
    # 包装为 Tool
    my_tool = Tool(
        name="build_recheck_tool",
        func=recheck_lecl_chain.invoke,  
        description="使用该链判断给出药方是否合适"
    )
    agent = create_react_agent(
        model=llm,
        tools=[my_tool],
        prompt="你需要根据输入的文本，调用工具进行操作。工具预期的输入格式为'recipe':'...'和'structured_summary':'...'以及'diagnosis_chat_history':'...'的字典。判断工具的返回结果是否正确，如果正确即刻返回结果。" # 可选，自定义 ReAct Prompt
    )
    return recheck_lecl_chain

def get_all_agent():
    return {key: fn() for key, fn in registry.items()}

def stream_graph_updates(user_input: str,agent):
    for event in agent.stream({"messages": [{"role": "user", "content":user_input}]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)


if __name__=="__main__":
    long_text =""""
    第一次就诊（2024.10.28）
医生：您好，哪里不舒服？
患者：医生，我发烧3天了，体温38.5℃，还有鼻塞、喉咙痛。
医生：有咳嗽咳痰吗？最近接触过流感病人吗？
患者：干咳没痰，同事上周也有感冒的。
医生：检查咽部充血，扁桃体不大。血常规显示病毒感染。开退烧药和抗病毒口服液，多休息多喝水。
患者：需要请假吗？
医生：建议休息3天，发热超过39℃或症状加重随时复诊。
第二次就诊（2024.11.18，间隔3周）
医生：上次用药后好些了吗？
患者：退烧了但一直流黄鼻涕，这两天又发烧到37.8℃。
医生：听诊肺部无啰音，鼻窦有压痛。可能合并细菌感染，加用抗生素，用生理盐水洗鼻。
患者：会传染给孩子吗？
医生：戴口罩避免密切接触，孩子出现症状要及时检查。
第三次就诊（2024.12.16，间隔4周）
医生：怎么又发烧了？
患者：这次是低烧37.5℃，但头痛乏力，咳嗽带黄痰。
医生：查体咽红，肺部呼吸音粗。建议拍胸片排除肺炎，近期免疫力可能较差。
患者：为什么反复感冒？
医生：需要查免疫功能和维生素D水平，平时注意锻炼和营养均衡。
    """
    
    structured_summary = """"
    第一次就诊（2024.10.28）

主诉：发烧3天（38.5℃）、鼻塞、喉咙痛，干咳无痰。

病史：同事近期感冒。

查体与辅助检查：咽部充血，扁桃体正常，血常规提示病毒感染。

处理：开退烧药和抗病毒药，建议多休息多喝水，休息3天，发热加重随诊。

第二次就诊（2024.11.18）

病情变化：退烧后持续黄鼻涕，两天前再次发烧至37.8℃。

查体：肺部无啰音，鼻窦压痛。

考虑：病毒感染后继发细菌感染。

处理：加用抗生素，生理盐水洗鼻。提醒戴口罩，注意避免传染孩子。

第三次就诊（2024.12.16）

病情变化：低烧（37.5℃），头痛、乏力，咳嗽黄痰。

查体：咽红，肺部呼吸音粗。

处理：建议拍胸片排除肺炎。


    """
    
    # 关键在于对agent的工具的描述和外部agent提示词的构造以及输入的构造
    # agent = build_history_agent()
    # res = agent.invoke({"messages":[
    #     {"role":"user","content":"病人病历是："+long_text},
    #     {"role":"user","content":"病人总结病史："+"病人没有总结病史。"},
    # ]})
    # print("agent sout is ",res["messages"][-1].content)
    
    agent = build_parse_agent()
    # res = agent.invoke({"messages":[
    #     {"role":"user","content":"病人症状是：头疼脑热,汗流浃背"},
    #     {"role":"user","content":"病人总结病史："+structured_summary},
    # ]})
    res = agent.invoke({"symptom":"有点不舒服","structured_summary":structured_summary})
    print("agent sout is ",res) 
    # while True:
    #     try:
    #         user_input = input("User: ")
    #         if user_input.lower() in ["quit", "exit", "q"]:
    #             print("Goodbye!")
    #             break
    #         stream_graph_updates(user_input,agent)
                
    #     except Exception as e:
    #         # fallback if input() is not available
            
    #         print(e)
            
    #         break
    